# 此模块负责下载器的文档处理逻辑

import re
from bs4 import BeautifulSoup
from bs4 import UnicodeDammit


class Document:

    def __init__(self, document):
        if type(document) is str:
            self.__document = document
        else:
            dammit = UnicodeDammit(document)
            encoding = dammit.original_encoding
            self.__document = document.decode(encoding=encoding)

        self.__bs = BeautifulSoup(self.__document, 'html.parser')
    
    def get_link_all(self, separate=False):
        """
        从 <link> <script> <image> 标签中提取链接
        :separate: 如果该参数为True,返回一个把链接分类后的字典 
        """
        tags_link = self.__bs.find_all('link', href=True)
        tags_script = self.__bs.find_all('script', src=True)
        tags_img = self.__bs.find_all('img', src=True)

        link_js = []
        for script in tags_script:
            link_js.append(script['src'])

        link_img = []
        for img in tags_img:
            link_img.append(img['src'])

        if separate:
            link_css = []
            link_ico = []
            for link in tags_link:
                if '.css' in link['href']:
                    link_css.append(link['href'])
                elif '.ico' in link['href']:
                    link_ico.append(link['href'])
                else:
                    picture_format = ['.gif', '.png', '.webp', 'jpg']
                    for pf in picture_format: 
                        link_img.append(link['href'])
            links = {}
            links['js'] = link_js
            links['image'] = link_img
            links['css'] = link_css
            links['ico'] = link_ico
        else:
            links = []
            links += link_js
            links += link_img
            for link in tags_link:
                links.append(link['href'])
        
        return links

    def replace(self, old, new) -> bool:
        """
        替换文档中的某部分内容
        :old: 旧内容
        :new: 替换后的新内容
        """
        document = self.__document.replace(old, new)
        if document == self.__document:
            return False
        else:
            self.__document = document
            return True

    def get_tag_str(self, tag_name):
        tag = self.__bs.find(tag_name)
        return tag.string

    def prettify(self) -> str:
        """
        输出文档
        """
        return self.__document